package io.kwdb.sdk;

import io.grpc.CallOptions;
import io.grpc.Channel;
import io.grpc.ClientCall;
import io.grpc.ClientInterceptor;
import io.grpc.ManagedChannel;
import io.grpc.ManagedChannelBuilder;
import io.grpc.Metadata;
import io.grpc.MethodDescriptor;
import io.grpc.StatusRuntimeException;
import io.grpc.netty.shaded.io.grpc.netty.GrpcSslContexts;
import io.grpc.netty.shaded.io.grpc.netty.NettyChannelBuilder;
import io.grpc.netty.shaded.io.netty.handler.ssl.SslContext;
import io.kwdb.proto.Ingest;
import io.kwdb.proto.IngestServiceGrpc;
import io.kwdb.proto.QueryServiceGrpc;
import io.kwdb.sdk.internal.*;
import java.io.File;
import java.util.Base64;
import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.TimeUnit;
import java.util.logging.Level;
import java.util.logging.Logger;

/**
 * gRPC implementation of the KwdbClient interface. This class manages the connection to the KWDB
 * server via gRPC and provides methods for writing data and executing queries. M3: Enhanced with
 * batching, retry, metrics, and security features.
 */
public class GrpcKwdbClient implements KwdbClient {

  private static final Logger logger = Logger.getLogger(GrpcKwdbClient.class.getName());

  private final Config config;
  private final ManagedChannel channel;
  private final IngestServiceGrpc.IngestServiceBlockingStub writeStub;
  private final QueryServiceGrpc.QueryServiceBlockingStub queryStub;
  private final RetryPolicy retryPolicy;
  private final Metrics metrics;
  private final Batcher<Point> batcher;
  private volatile boolean closed = false;

  /**
   * Create a new GrpcKwdbClient with the given configuration.
   *
   * @param config The client configuration
   */
  public GrpcKwdbClient(Config config) {
    Validators.validateConfig(config);
    this.config = config;
    this.metrics = Metrics.getInstance();

    // M3: Build the gRPC channel with security and compression support
    this.channel = buildChannel(config);

    // M3: Create stubs with auth interceptor if needed
    if (config.getAuthToken() != null || config.getUsername() != null) {
      ClientInterceptor authInterceptor = new AuthInterceptor(config);
      this.writeStub = IngestServiceGrpc.newBlockingStub(channel).withInterceptors(authInterceptor);
      this.queryStub = QueryServiceGrpc.newBlockingStub(channel).withInterceptors(authInterceptor);
    } else {
      this.writeStub = IngestServiceGrpc.newBlockingStub(channel);
      this.queryStub = QueryServiceGrpc.newBlockingStub(channel);
    }

    // M3: Initialize retry policy
    this.retryPolicy =
        new RetryPolicy(config.isEnableRetry(), config.getMaxRetries(), config.getBaseBackoffMs());

    // M3: Initialize batcher if enabled
    if (config.getMaxBatchSize() > 0) {
      this.batcher =
          new Batcher<>(
              config.getMaxBatchSize(),
              config.getLingerMs(),
              config.getMaxQueueSize(),
              config.getBackpressure(),
              this::submitBatch);
      logger.info("gRPC client initialized with batcher enabled");
    } else {
      this.batcher = null;
      logger.info("gRPC client initialized with direct write mode");
    }

    logger.info(
        String.format("GrpcKwdbClient created for %s:%d", config.getHost(), config.getPort()));
  }

  /** Build gRPC channel with TLS and compression configuration. */
  private ManagedChannel buildChannel(Config config) {
    try {
      if (config.isTls()) {
        // M3: Build secure channel with TLS
        NettyChannelBuilder channelBuilder =
            NettyChannelBuilder.forAddress(config.getHost(), config.getPort());

        if (config.getCaCertPath() != null && !config.getCaCertPath().isEmpty()) {
          // Use custom CA certificate
          File caCert = new File(config.getCaCertPath());
          SslContext sslContext = GrpcSslContexts.forClient().trustManager(caCert).build();
          channelBuilder.sslContext(sslContext);
        } else {
          // Use default system trust store
          channelBuilder.useTransportSecurity();
        }

        if (config.isGzip()) {
          channelBuilder.enableRetry().maxRetryAttempts(config.getMaxRetries());
        }

        return channelBuilder.build();
      } else {
        // Plain text channel
        ManagedChannelBuilder<?> channelBuilder =
            ManagedChannelBuilder.forAddress(config.getHost(), config.getPort()).usePlaintext();

        if (config.isGzip()) {
          channelBuilder.enableRetry().maxRetryAttempts(config.getMaxRetries());
        }

        return channelBuilder.build();
      }
    } catch (Exception e) {
      logger.log(Level.WARNING, "Failed to build secure channel, falling back to plaintext", e);
      return ManagedChannelBuilder.forAddress(config.getHost(), config.getPort())
          .usePlaintext()
          .build();
    }
  }

  @Override
  public Status write(Point point) {
    if (closed) {
      return Status.error(Status.Code.FAILED_PRECONDITION, "Client is closed");
    }

    if (point == null) {
      return Status.invalidArgument("Point cannot be null");
    }

    // M3: Use batcher if enabled
    if (batcher != null) {
      boolean accepted = batcher.offer(point, config.getTimeoutMillis());
      if (accepted) {
        return Status.ok();
      } else {
        metrics.incrementDropped(1);
        return Status.resourceExhausted("Failed to enqueue point");
      }
    }

    // Direct write with retry
    return writeDirectWithRetry(point);
  }

  /** Write a single point directly with retry logic. */
  private Status writeDirectWithRetry(Point point) {
    long startTime = System.currentTimeMillis();

    Status status =
        retryPolicy.execute(
            () -> {
              try {
                Validators.validatePoint(point);

                // Convert point to protobuf request
                Ingest.WriteRequest request = ProtoMappers.toWriteRequest(point);

                // Send request with timeout
                Ingest.WriteResponse response =
                    writeStub
                        .withDeadlineAfter(config.getTimeoutMillis(), TimeUnit.MILLISECONDS)
                        .writeBatch(request);

                // Check response status (code 0 = success)
                if (response.getCode() == 0) {
                  return Status.ok();
                } else {
                  String errorMsg =
                      response.getMessage() != null && !response.getMessage().isEmpty()
                          ? response.getMessage()
                          : "Write failed";
                  return Status.error(Status.Code.fromValue(response.getCode()), errorMsg);
                }

              } catch (IllegalArgumentException e) {
                return Status.invalidArgument(e.getMessage());
              } catch (StatusRuntimeException e) {
                return Status.fromGrpc(e);
              } catch (Exception e) {
                return Status.internal("Unexpected error: " + e.getMessage());
              }
            });

    long elapsed = System.currentTimeMillis() - startTime;
    metrics.observeWriteLatency(elapsed);

    if (status.isOk()) {
      metrics.incrementWritesOk(1);
    } else {
      metrics.incrementWritesFailed(1, status.getCode().getValue());
    }

    return status;
  }

  @Override
  public Status batchWrite(List<Point> batch) {
    if (closed) {
      return Status.error(Status.Code.FAILED_PRECONDITION, "Client is closed");
    }

    if (batch == null || batch.isEmpty()) {
      return Status.invalidArgument("Batch cannot be null or empty");
    }

    // M3: Direct batch write bypasses batcher
    return submitBatchDirectly(batch);
  }

  /** Submit a batch of points directly (used by batchWrite and batcher). */
  private Status submitBatchDirectly(List<Point> batch) {
    long startTime = System.currentTimeMillis();

    Status status =
        retryPolicy.execute(
            () -> {
              try {
                Validators.validateBatch(batch);

                // Convert batch to protobuf request
                Ingest.WriteRequest request = ProtoMappers.toWriteRequest(batch);

                // Send request with timeout
                Ingest.WriteResponse response =
                    writeStub
                        .withDeadlineAfter(config.getTimeoutMillis(), TimeUnit.MILLISECONDS)
                        .writeBatch(request);

                // Check response status (code 0 = success)
                if (response.getCode() == 0) {
                  return Status.ok();
                } else {
                  String errorMsg =
                      response.getMessage() != null && !response.getMessage().isEmpty()
                          ? response.getMessage()
                          : "Batch write failed";
                  return Status.error(Status.Code.fromValue(response.getCode()), errorMsg);
                }

              } catch (IllegalArgumentException e) {
                return Status.invalidArgument(e.getMessage());
              } catch (StatusRuntimeException e) {
                return Status.fromGrpc(e);
              } catch (Exception e) {
                return Status.internal("Unexpected error: " + e.getMessage());
              }
            });

    long elapsed = System.currentTimeMillis() - startTime;
    metrics.observeBatchLatency(elapsed);
    metrics.observeBatchSize(batch.size());

    if (status.isOk()) {
      metrics.incrementWritesOk(batch.size());
    } else {
      metrics.incrementWritesFailed(batch.size(), status.getCode().getValue());
    }

    return status;
  }

  /** Callback for batcher to submit accumulated points. */
  private void submitBatch(List<Point> batch) {
    if (batch.isEmpty()) {
      return;
    }

    Status status = submitBatchDirectly(batch);
    if (!status.isOk()) {
      logger.log(Level.WARNING, "Batch submission failed: {0}", status.getMessage());
    }
  }

  @Override
  public QueryResult query(String sql) {
    if (closed) {
      return QueryResult.error(Status.error(Status.Code.FAILED_PRECONDITION, "Client is closed"));
    }

    try {
      Validators.validateSql(sql);

      // Create query request
      Ingest.QueryReq queryRequest = ProtoMappers.toQueryReq(sql);

      // Send query with timeout
      Ingest.QueryResp response =
          queryStub
              .withDeadlineAfter(config.getTimeoutMillis(), TimeUnit.MILLISECONDS)
              .query(queryRequest);

      // Convert response to QueryResult
      QueryResult result = ProtoMappers.fromQueryResp(response);

      logger.fine(String.format("Query executed, returned %d rows", result.getRowCount()));
      return result;

    } catch (IllegalArgumentException e) {
      logger.warning("Invalid SQL: " + e.getMessage());
      return QueryResult.error(Status.invalidArgument(e.getMessage()));
    } catch (StatusRuntimeException e) {
      logger.log(Level.WARNING, "gRPC error during query", e);
      return QueryResult.error(Status.fromGrpc(e));
    } catch (Exception e) {
      logger.log(Level.SEVERE, "Unexpected error during query", e);
      return QueryResult.error(Status.internal("Unexpected error: " + e.getMessage()));
    }
  }

  @Override
  public CompletableFuture<Status> writeAsync(Point point) {
    if (closed) {
      return CompletableFuture.completedFuture(
          Status.error(Status.Code.FAILED_PRECONDITION, "Client is closed"));
    }

    return CompletableFuture.supplyAsync(() -> write(point));
  }

  @Override
  public void flush() {
    if (batcher != null && !closed) {
      batcher.flush();
    }
  }

  @Override
  public void close() throws Exception {
    if (closed) {
      return;
    }

    closed = true;
    logger.info("Closing gRPC client...");

    // M3: Close batcher first (will flush pending items)
    if (batcher != null) {
      try {
        batcher.close();
      } catch (Exception e) {
        logger.log(Level.WARNING, "Error closing batcher", e);
      }
    }

    try {
      // Shutdown the channel gracefully
      channel.shutdown();

      // Wait for termination with timeout
      if (!channel.awaitTermination(5, TimeUnit.SECONDS)) {
        logger.warning("Channel did not terminate in 5 seconds, forcing shutdown");
        channel.shutdownNow();

        if (!channel.awaitTermination(5, TimeUnit.SECONDS)) {
          logger.severe("Channel did not terminate after forced shutdown");
        }
      }

      // Log final metrics
      Metrics.MetricsSnapshot snapshot = metrics.getSnapshot();
      logger.info("Final metrics: " + snapshot);
      logger.info("GrpcKwdbClient closed successfully");

    } catch (InterruptedException e) {
      Thread.currentThread().interrupt();
      logger.log(Level.WARNING, "Interrupted during channel shutdown", e);
      channel.shutdownNow();
      throw new Exception("Interrupted during close", e);
    }
  }

  @Override
  public boolean isClosed() {
    return closed;
  }

  /**
   * Get the configuration used by this client. This is mainly for testing and debugging purposes.
   *
   * @return The client configuration
   */
  public Config getConfig() {
    return config;
  }

  /**
   * Check if the underlying gRPC channel is still active. This is mainly for testing and debugging
   * purposes.
   *
   * @return true if the channel is not shutdown or terminated
   */
  public boolean isChannelActive() {
    return !channel.isShutdown() && !channel.isTerminated();
  }

  /** gRPC interceptor for authentication headers. */
  private static class AuthInterceptor implements ClientInterceptor {
    private static final Metadata.Key<String> AUTH_HEADER =
        Metadata.Key.of("authorization", Metadata.ASCII_STRING_MARSHALLER);

    private final String authValue;

    public AuthInterceptor(Config config) {
      if (config.getAuthToken() != null && !config.getAuthToken().isEmpty()) {
        this.authValue = "Bearer " + config.getAuthToken();
      } else if (config.getUsername() != null && config.getPassword() != null) {
        String credentials = config.getUsername() + ":" + config.getPassword();
        String encoded =
            Base64.getEncoder()
                .encodeToString(credentials.getBytes(java.nio.charset.StandardCharsets.UTF_8));
        this.authValue = "Basic " + encoded;
      } else {
        this.authValue = null;
      }
    }

    @Override
    public <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall(
        MethodDescriptor<ReqT, RespT> method, CallOptions callOptions, Channel next) {

      return new ClientCall<ReqT, RespT>() {
        private ClientCall<ReqT, RespT> call;

        @Override
        public void start(Listener<RespT> responseListener, Metadata headers) {
          if (authValue != null) {
            headers.put(AUTH_HEADER, authValue);
          }
          call = next.newCall(method, callOptions);
          call.start(responseListener, headers);
        }

        @Override
        public void request(int numMessages) {
          call.request(numMessages);
        }

        @Override
        public void cancel(String message, Throwable cause) {
          call.cancel(message, cause);
        }

        @Override
        public void halfClose() {
          call.halfClose();
        }

        @Override
        public void sendMessage(ReqT message) {
          call.sendMessage(message);
        }
      };
    }
  }
}
